'''
train，test，validation 

parser.add_argument('-j', '--workers', default=16, type=int, metavar='N',
                    help='number of data loading workers (default: 16)')

'''
import os
import pandas as pds
import argparse
import xml.etree.ElementTree as Etree
import math
import random
from  imageRender import xmlnodeToJson,XMLToJson

parse=argparse.ArgumentParser(description="split dataset")

parse.add_argument("--root",default="/media/gis/data/jupyterlabhub/gitcode/hrx/dataset",type=str,help="the root dir of the dataset")
parse.add_argument("--labeldir",default="biaoqian",type=str,help="the dir name of label")
parse.add_argument("--imagedir",default="shujuji",type=str,help="the dir name of image")
parse.add_argument("--train",default=0.7,type=float,help="the number of train is divided by all images")
parse.add_argument("--test",default=0.1,type=float,help="the number of test is divided by all images")
parse.add_argument("--vail",default=0.2,type=float,help="the number of vail is divided by all images")

args=parse.parse_args() #参数解析
rootdir=args.root
anndir=args.labeldir
imgdir=args.imagedir

# 开始分割文件

splitcache=[]


# 确定可以参与分割的文件情况
annlist=os.listdir(os.path.join(rootdir,anndir))
imglist=os.listdir(os.path.join(rootdir,imgdir))
for an in annlist:
    antxtpath=os.path.join(rootdir,anndir,an)
    conxml=Etree.parse(antxtpath)
    timg_path=conxml.find("path").text
    tpath=os.path.join(rootdir,timg_path)
    if (os.path.exists(tpath)):
        splitcache.append(an)

csvname="train_test_vail.csv"
csvpath=os.path.join(rootdir,csvname)

result={"id":[],"annotation":[],"images":[],"sig":[],"class":[]}
i=1
# 进行文件拆分开始
for an in splitcache:
    antxtpath=os.path.join(rootdir,anndir,an)
    conxml=Etree.parse(antxtpath)
    timg_path=conxml.find("path").text
    tpath=os.path.join(rootdir,timg_path)
    # 训练分图 
    result["id"].append(i)
    result["annotation"].append(antxtpath)
    result["images"].append(tpath)
    if (os.path.exists(tpath)):
        sig=random.random()
        if sig<args.train:
            result["sig"].append("train")
        elif sig<args.train+args.test:
            result["sig"].append("test")
        elif sig<args.train+args.test+args.vail:
            result["sig"].append("vail")
    # 图片分类
    objlist=conxml.findall("object")
    D1Ct=0
    for obj in objlist:
        clsname=obj.find("name")
        if clsname.text=="D1":
            D1Ct=1
    if D1Ct==1:
        result["class"].append(1)
    else:
        result["class"].append(0)
    i=i+1

result = pds.DataFrame(result)
result.to_csv(csvpath,encoding="utf-8")